#include "SphericalHarmonics.h"

//******************************************************************
//FUNCTION:
int degFromSh(int vNumBases) {
	switch (vNumBases) {
	case 1:
		return 0;
	case 4:
		return 1;
	case 9:
		return 2;
	case 16:
		return 3;
	default:
		return 4;
	}
}
namespace sphericalHarmonics {
	const double C0 = 0.28209479177387814;
}

//******************************************************************
//FUNCTION:
torch::Tensor rgb2sh(const torch::Tensor& vRgb) {
	//Author Comment: Converts from RGB values [0,1] to the 0th spherical harmonic coefficient
	return (vRgb - 0.5) / sphericalHarmonics::C0;
}

//******************************************************************
//FUNCTION:
torch::Tensor sh2rgb(const torch::Tensor& vSH) {
	//Author Comment: Converts from 0th spherical harmonic coefficients to RGB values [0,1]
	return torch::clamp((vSH * sphericalHarmonics::C0) + 0.5, 0.0f, 1.0f);
}

//******************************************************************
//FUNCTION:
torch::Tensor CSphericalHarmonics::forward(AutogradContext* vCtx,
	int vDegreesToUse,
	torch::Tensor vViewDirs,
	torch::Tensor vCoeffs) {
	long long NumPoints = vCoeffs.size(0);
	int Degree = degFromSh(vCoeffs.size(-2));

	vCtx->saved_data["degreesToUse"] = vDegreesToUse;
	vCtx->saved_data["degree"] = Degree;

	vCtx->save_for_backward({ vViewDirs });

	return compute_sh_forward_tensor(NumPoints, Degree, vDegreesToUse, vViewDirs, vCoeffs);
}

//******************************************************************
//FUNCTION:
tensor_list CSphericalHarmonics::backward(AutogradContext* vCtx, tensor_list vGradOutputs) {
	torch::Tensor Colors = vGradOutputs[0];
	int DegreesToUse = vCtx->saved_data["degreesToUse"].toInt();
	int Degree = vCtx->saved_data["degree"].toInt();
	variable_list Saved = vCtx->get_saved_variables();

	torch::Tensor ViewDirs = Saved[0];
	long long NumPoints = Colors.size(0);
	torch::Tensor None;

	return {
		None,
		None,
		compute_sh_backward_tensor(NumPoints, Degree, DegreesToUse, ViewDirs, Colors)
	};
}